package cc.shacocloud.mirage.maven.plugin;

import org.apache.commons.compress.archivers.jar.JarArchiveEntry;
import org.apache.commons.compress.archivers.jar.JarArchiveOutputStream;
import org.apache.commons.compress.archivers.zip.UnixStat;
import org.apache.maven.artifact.Artifact;
import org.jetbrains.annotations.Contract;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.attribute.FileTime;
import java.util.*;
import java.util.function.Supplier;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;
import java.util.jar.JarInputStream;
import java.util.jar.Manifest;
import java.util.zip.CRC32;
import java.util.zip.ZipEntry;
import java.util.zip.ZipException;

/**
 * 向 JAR 写入内容
 *
 * @author 思追(shaco)
 */
public class JarWriter implements AutoCloseable {
    
    private static final int UNIX_FILE_MODE = UnixStat.FILE_FLAG | UnixStat.DEFAULT_FILE_PERM;
    private static final int UNIX_DIR_MODE = UnixStat.DIR_FLAG | UnixStat.DEFAULT_DIR_PERM;
    
    private final JarArchiveOutputStream jarOutputStream;
    
    private final FileTime lastModifiedTime;
    
    private final Set<String> writtenEntries = new HashSet<>();
    
    public JarWriter(File file) throws IOException {
        this(file, null);
    }
    
    public JarWriter(File file, FileTime lastModifiedTime) throws IOException {
        this.jarOutputStream = new JarArchiveOutputStream(new FileOutputStream(file));
        this.jarOutputStream.setEncoding("UTF-8");
        this.lastModifiedTime = lastModifiedTime;
    }
    
    /**
     * 写入 {@link Artifact} 依赖，返回对应的 {@link JarArchiveEntry}
     */
    @Nullable
    public JarArchiveEntry writeArtifactLib(@NotNull ArtifactLib artifactLib, @NotNull JarArchiveEntryTransformer entryTransformer) throws Exception {
        JarArchiveEntry entry = new JarArchiveEntry(artifactLib.getName());
        entry.setTime(getArtifactLibTime(artifactLib));
        new CrcAndSizeHolder(artifactLib::openStream).setupStoredEntry(entry);
        
        try (InputStream inputStream = artifactLib.openStream()) {
            return writeEntry(entry, inputStream, entryTransformer);
        }
    }
    
    /**
     * 复制jar文件的内容，使用 {@link JarArchiveEntryTransformer} 对于条目进行转换
     */
    public void copyByJarFile(@NotNull JarFile sourceJar,
                              @NotNull JarArchiveEntryTransformer entryTransformer) throws Exception {
        Enumeration<JarEntry> entries = sourceJar.entries();
        
        while (entries.hasMoreElements()) {
            JarEntry jarEntry = entries.nextElement();
            JarArchiveEntry jarArchiveEntry = new JarArchiveEntry(jarEntry);
            writeEntry(sourceJar, jarArchiveEntry, entryTransformer);
        }
    }
    
    /**
     * 写 Manifest
     */
    public void writeManifest(@NotNull Manifest manifest) throws IOException {
        JarArchiveEntry entry = new JarArchiveEntry("META-INF/MANIFEST.MF");
        writeEntry(entry, manifest::write);
    }
    
    /**
     * 使用输入流写入 jar 条目
     */
    @Nullable
    public JarArchiveEntry writeEntry(@NotNull JarArchiveEntry entry,
                                      @NotNull InputStream inputStream,
                                      @NotNull JarArchiveEntryTransformer entryTransformer) throws IOException {
        JarEntryWriter entryWriter = new InputStreamJarEntryWriter(inputStream);
        JarArchiveEntry transformedEntry = entryTransformer.transform(entry);
        if (transformedEntry != null) {
            writeEntry(transformedEntry, entryWriter);
            
            return transformedEntry;
        }
        
        return null;
    }
    
    /**
     * 写入 jar 条目
     */
    public void writeEntry(@NotNull JarArchiveEntry entry, @Nullable JarEntryWriter entryWriter) throws IOException {
        String name = entry.getName();
        if (this.writtenEntries.add(name)) {
            writeParentDirectoryEntries(name);
            entry.setUnixMode(name.endsWith("/") ? UNIX_DIR_MODE : UNIX_FILE_MODE);
            entry.getGeneralPurposeBit().useUTF8ForNames(true);
            if (!entry.isDirectory() && entry.getSize() == -1) {
                
                if (Objects.nonNull(entryWriter)) {
                    entryWriter = new CalculateSizeJarEntryWriter(entryWriter);
                    entry.setSize(entryWriter.size());
                } else {
                    entry.setSize(0);
                }
            }
            
            writeToArchive(entry, entryWriter);
        }
    }
    
    private void writeParentDirectoryEntries(@NotNull String name) throws IOException {
        String parent = name.endsWith("/") ? name.substring(0, name.length() - 1) : name;
        while (parent.lastIndexOf('/') != -1) {
            parent = parent.substring(0, parent.lastIndexOf('/'));
            if (!parent.isEmpty()) {
                writeEntry(new JarArchiveEntry(parent + "/"), null);
            }
        }
    }
    
    private void writeToArchive(ZipEntry entry, @Nullable JarEntryWriter entryWriter) throws IOException {
        JarArchiveEntry jarEntry = asJarArchiveEntry(entry);
        if (this.lastModifiedTime != null) {
            long time = this.lastModifiedTime.toMillis();
            FileTime fileTime = FileTime.fromMillis(time - TimeZone.getDefault().getOffset(time));
            jarEntry.setTime(fileTime);
        }
        this.jarOutputStream.putArchiveEntry(jarEntry);
        if (entryWriter != null) {
            entryWriter.write(this.jarOutputStream);
        }
        this.jarOutputStream.closeArchiveEntry();
    }
    
    private void writeEntry(@NotNull JarFile jarFile,
                            @NotNull JarArchiveEntry entry,
                            @NotNull JarArchiveEntryTransformer entryTransformer) throws IOException {
        setUpEntry(jarFile, entry);
        try (ZipHeaderPeekInputStream inputStream = new ZipHeaderPeekInputStream(jarFile.getInputStream(entry))) {
            writeEntry(entry, inputStream, entryTransformer);
        }
    }
    
    private void setUpEntry(@NotNull JarFile jarFile, @NotNull JarArchiveEntry entry) throws IOException {
        try (ZipHeaderPeekInputStream inputStream = new ZipHeaderPeekInputStream(jarFile.getInputStream(entry))) {
            if (inputStream.hasZipHeader() && entry.getMethod() != ZipEntry.STORED) {
                new CrcAndSizeHolder(inputStream).setupStoredEntry(entry);
            } else {
                entry.setCompressedSize(-1);
            }
        }
    }
    
    @NotNull
    @Contract("null -> new")
    private JarArchiveEntry asJarArchiveEntry(ZipEntry entry) throws ZipException {
        if (entry instanceof JarArchiveEntry) {
            return (JarArchiveEntry) entry;
        }
        return new JarArchiveEntry(entry);
    }
    
    private long getArtifactLibTime(ArtifactLib artifactLib) {
        try {
            try (JarInputStream jarStream = new JarInputStream(artifactLib.openStream())) {
                JarEntry entry = jarStream.getNextJarEntry();
                while (entry != null) {
                    if (!entry.isDirectory()) {
                        return entry.getTime();
                    }
                    entry = jarStream.getNextJarEntry();
                }
            }
        } catch (Exception ex) {
            // 忽略并仅使用库时间戳
        }
        return artifactLib.getLastModified();
    }
    
    @Override
    public void close() throws IOException {
        this.jarOutputStream.close();
    }
    
    private static class CrcAndSizeHolder {
        
        private static final int BUFFER_SIZE = 32 * 1024;
        
        private final CRC32 crc = new CRC32();
        
        private long size;
        
        CrcAndSizeHolder(@NotNull Supplier<InputStream> supplier) throws IOException {
            try (InputStream inputStream = supplier.get()) {
                load(inputStream);
            }
        }
        
        CrcAndSizeHolder(InputStream inputStream) throws IOException {
            load(inputStream);
        }
        
        private void load(@NotNull InputStream inputStream) throws IOException {
            byte[] buffer = new byte[BUFFER_SIZE];
            int bytesRead;
            while ((bytesRead = inputStream.read(buffer)) != -1) {
                this.crc.update(buffer, 0, bytesRead);
                this.size += bytesRead;
            }
        }
        
        void setupStoredEntry(@NotNull JarArchiveEntry entry) {
            entry.setSize(this.size);
            entry.setCompressedSize(this.size);
            entry.setCrc(this.crc.getValue());
            entry.setMethod(ZipEntry.STORED);
        }
        
    }
    
}
